#!/usr/bin/env python

import gtk
import gtk.glade
import gtk.gtkgl

from OpenGL.GL import *
from OpenGL.GLU import *
from OpenGL.raw import GL
from OpenGL.arrays import ArrayDatatype as ADT

from PIL import Image
from scipy import array, cross, sqrt, square, cos, sin, pi, float32, radians, vdot


class SceneNode:
	def __init__(self, rotation = (0.0, 0.0, 0.0), 
				       position = (0.0, 0.0, 0.0)):
		self.subnodes = {}
		self.position = position
		self.rotation = rotation
		self.matrix = None
		self.parent = None
	
	def push_matrix(self):
		glPushMatrix()
		if self.matrix:
			glMultMatrixd(self.matrix)
		else:
			glTranslatef(*self.position)
			glRotatef(self.rotation[2], 0.0, 0.0, 1.0)
			glRotatef(self.rotation[1], 0.0, 1.0, 0.0)
			glRotatef(self.rotation[0], 1.0, 0.0, 0.0)

	def pop_matrix(self):
		glPopMatrix()
	
	def draw(self):
		self.push_matrix()
		for keys in self.subnodes:
			self.subnodes[keys].draw()
		self.pop_matrix()
	
	def add(self, name, subnode):
		subnode.parent = self
		self.subnodes[name] = subnode
	
	def __getitem__(self, key):
		return self.subnodes[key]
	
	def set_position_in_meters(self, pos):
		self.position = (pos[0]*1000, pos[1]*1000+200, pos[2]*1000)

	def set_rotation_in_radians(self, pos):
		self.rotation = (pos[0]*57.2957795, pos[1]*57.2957795, pos[2]*57.2957795)
	
	def set_matrix(self, mtrx):
		self.matrix = mtrx
	
	def make_origo(self):
		origo_x = ScenePart()
		origo_x.make_arrow(3*0.001, 1*0.001, 18*0.001, color=(1.0,0.0,0.0,1.0))
		origo_x.rotation = (0, 0, -90)
		origo_x.position = (20*0.001, 0, 0)
		self.add("ox", origo_x)		
		origo_y = ScenePart()
		origo_y.make_arrow(3*0.001, 1*0.001, 18*0.001, color=(0.0,1.0,0.0,1.0))
		origo_y.rotation = (90, 0, 0)
		origo_y.position = (0, 0, 20*0.001)
		self.add("oy", origo_y)		
		origo_z = ScenePart()
		origo_z.make_arrow(3*0.001, 1*0.001, 18*0.001, color=(0.0,0.0,1.0,1.0))
		origo_z.position = (0, 20*0.001, 0)
		self.add("oz", origo_z)
	
class ScenePart(SceneNode):
	def __init__(self, filename = None, 
				       color = (0.2, 0.2, 0.2, 1.0), 
				       rotation = (0.0, 0.0, 0.0), 
				       position = (0.0, 0.0, 0.0)):
		SceneNode.__init__(self, rotation = rotation, position = position)
		self.vertices = None
		self.normals = None
		self.vertex_size = None
		self.vertex_type = None
		self.color = color
		self.color_specular = (0.5, 0.5, 0.5, 1.0)
		if filename:
			self.load_file(filename)
		
	def draw(self):
		self.push_matrix()		
		p = glVertexPointer(self.vertex_size, GL_FLOAT, 0, self.vertices)
		p2 = glNormalPointer(GL_FLOAT, 0, self.normals)		
		glEnableClientState(GL_VERTEX_ARRAY)
		glEnableClientState(GL_NORMAL_ARRAY)
		glMaterialfv(GL_FRONT, GL_AMBIENT_AND_DIFFUSE, self.color)
		glMateriali(GL_FRONT, GL_SHININESS, 70)
		glMaterialfv(GL_FRONT, GL_SPECULAR, self.color_specular)
		glDrawArrays(self.vertex_type, 0, len(self.vertices)/self.vertex_size)		
		self.pop_matrix()
	
	def make_box(self, size, position = (0.0, 0.0, 0.0)):
		self.position = position
		self.quads = []
		self.triangles = []
		x, y, z = size[0]/2.0, size[1]/2.0, size[2]/2.0
		def add_quad(quad, normal):
			self.quads.append(Quad(quad, normal))
		add_quad([[-x,  y,  z], [ x,  y,  z], [ x, -y,  z], [-x, -y,  z]], [ 0, 0, 1])
		add_quad([[ x,  y,  z], [ x,  y, -z], [ x, -y, -z], [ x, -y,  z]], [ 1, 0, 0])
		add_quad([[ x,  y, -z], [-x,  y, -z], [-x, -y, -z], [ x, -y, -z]], [ 0, 0,-1])
		add_quad([[-x,  y, -z], [-x,  y,  z], [-x, -y,  z], [-x, -y, -z]], [-1, 0, 0])
		add_quad([[-x,  y,  z], [-x,  y, -z], [ x,  y, -z], [ x,  y,  z]], [ 0, 1, 0])
		add_quad([[-x, -y,  z], [ x, -y,  z], [ x, -y, -z], [-x, -y, -z]], [ 0,-1, 0])
	
	def make_cylinder(self, radius, height, segments = 20, color = (0.5, 0.5, 0.5, 1.0)):
		self.triangles = []
		self.quads = []
		self.color = color
		def add_quad(quad, normal):
			self.quads.append(Quad(quad, normal))
		def add_triangle(triangle, normal):
			self.triangles.append(Triangle(triangle, normal))
		prev_x, prev_z = (1.0, 0.0)
		for angle in range(1, segments+1):
			rad = float(angle) / segments * pi * 2
			x, z = cos(rad), sin(rad)
			add_quad([[x,  0,  z],[prev_x, 0, prev_z],[prev_x, height, prev_z], [x,  height,  z]], [x,0,z])
			add_triangle([[x, 0, z], [prev_x, 0, prev_z], [0, 0, 0]], [0,-1,0])
			add_triangle([[x, height, z], [prev_x, height, prev_z], [0, height, 0]], [0,1,0])
			prev_x, prev_z = x, z
	
	def make_cone(self, radius, height, segments = 100, color = (0.5, 0.5, 0.5, 1.0)):
		vertices = []
		normals = []
		self.color = color
		for angle in range(0, segments):
			rad1 = float(angle) / segments * pi * 2
			rad2 = float(angle+0.5) / segments * pi * 2
			rad3 = float(angle+1.0) / segments * pi * 2
			x1, z1 = cos(rad1)*radius, sin(rad1)*radius
			x2, z2 = cos(rad2)*radius, sin(rad2)*radius
			x3, z3 = cos(rad3)*radius, sin(rad3)*radius
			v1, v2, v3 = array([x1, 0, z1]), array([x3, 0, z3]), array([0, height, 0])
			n = cross(v3-v1, v2-v1)
			n /= sqrt(vdot(n,n))
			vertices.extend([x3, 0, z3, x1, 0, z1, 0, height, 0, x3, 0, z3, x1, 0, z1, 0, 0, 0])
			normals.extend([n[0], n[1], n[2], n[0], n[1], n[2], n[0], n[1], n[2], 0, -1, 0, 0, -1, 0, 0, -1, 0])
		
		self.vertices = array(vertices, dtype=float32)
		self.normals = array(normals, dtype=float32)
		self.vertex_size = 3
		self.vertex_type = GL_TRIANGLES
	
	def make_arrow(self, r1, r2, length, segments = 50, color = (0.5, 0.5, 0.5, 1.0)):
		# r1 - arrowhead radius 4
		# r2 - cone radius      2
		def add_vertices(v1, v2, v3):
			n = cross(array(v3)-array(v1), array(v2)-array(v1))
			n /= sqrt(vdot(n,n))
			vertices.extend([v1[0], v1[1], v1[2], v2[0], v2[1], v2[2], v3[0], v3[1], v3[2]])
			normals.extend([n[0], n[1], n[2], n[0], n[1], n[2], n[0], n[1], n[2]])
		
		h = r1 * 3
		vertices = []
		normals = []
		self.color = color
		for angle in range(0, segments):
			rad1 = float(angle) / segments * pi * 2
			rad2 = float(angle+0.5) / segments * pi * 2
			rad3 = float(angle+1.0) / segments * pi * 2
			x1, z1 = cos(rad1), sin(rad1)
			x2, z2 = cos(rad2), sin(rad2)
			x3, z3 = cos(rad3), sin(rad3)
			add_vertices([x1*r1, -h, z1*r1], [x3*r1, -h, z3*r1], [0, 0, 0])
			add_vertices([x1*r1, -h, z1*r1], [x1*r2, -h, z1*r2], [x3*r1, -h, z3*r1])
			add_vertices([x3*r1, -h, z3*r1], [x1*r2, -h, z1*r2], [x3*r2, -h, z3*r2])
			add_vertices([x1*r2, -h, z1*r2], [x1*r2, -length, z1*r2], [x3*r2, -length, z3*r2])
			add_vertices([x1*r2, -h, z1*r2], [x3*r2, -length, z3*r2], [x3*r2, -h, z3*r2])
			add_vertices([x1*r2, -length, z1*r2], [0, -length, 0], [x3*r2, -length, z3*r2])
			
		self.vertices = array(vertices, dtype=float32)
		self.normals = array(normals, dtype=float32)
		self.vertex_size = 3
		self.vertex_type = GL_TRIANGLES
	
	def load_file(self, filename):
		vertices = []
		normals = []
		f = open(filename, 'r')
		while True:
			line = f.readline().split()
			if len(line) == 0:
				break
			if line[0] == "facet":
				for idx in range(0, 3):
					normals.extend([float(line[2]), float(line[3]), float(line[4])])
				line = f.readline().split()
				for vidx in range(0, 3):
					line = f.readline().split()
					vertices.extend([float(line[1])*0.001, float(line[2])*0.001, float(line[3])*0.001])
		f.close()
		self.vertices = array(vertices, dtype=float32)
		self.normals = array(normals, dtype=float32)
		self.vertex_size = 3
		self.vertex_type = GL_TRIANGLES


class Camera:
	def __init__(self):
		self.center = array([0.0, 180.0*0.001, 0.0])
		self.rotation = array([0.7, 0.0, 0.0])
		self.up = array([0.0, 1.0, 0.0])
		self.distance = 300.0*0.001
		
	def get_camera_pos(self):
		r = self.distance
		rot_x, rot_y = self.rotation[0], self.rotation[1]
		eye = array([-r * sin(rot_x) * cos(rot_y),
					  r * sin(rot_y),
					  r * cos(rot_x) * cos(rot_y)])
		eye += self.center
		return eye
		
	def set_gl_orientation(self):
		up = self.up
		center = self.center
		eye = self.get_camera_pos()
		
		glLightfv(GL_LIGHT1, GL_POSITION, (-1000,500,1000,1.0))
		gluLookAt(eye[0], eye[1], eye[2],
				  center[0], center[1], center[2], 
				  up[0], up[1], up[2])
	
	def rotate_rel(self, val):
		self.rotation[0] += val[0] * 0.01
		self.rotation[1] += val[1] * 0.01
		pi_2 = pi / 2.0
		if self.rotation[1] > pi_2:
			self.rotation[1] = pi_2
		elif self.rotation[1] < -pi_2:
			self.rotation[1] = -pi_2

	def translate_rel(self, val):
		norm = lambda x: sqrt(square(x).sum())
		eye = self.center - self.get_camera_pos()
		eye = eye / norm(eye)
		e1 = self.up
		e2 = cross(eye, e1)
		self.center += e1 * val[1] * 0.8*0.001
		self.center -= e2 * val[0] * 0.8*0.001

	def zoom_rel(self, val):
		self.distance *= val


class VisualiseWidget3D(gtk.DrawingArea, gtk.gtkgl.Widget):
	def __init__(self, connection = None, save_video = False):
		self.connection = connection
		self.running = False
		self.save_video = save_video
		gtk.gdk.threads_init()
		gtk.DrawingArea.__init__(self)
		self.set_size_request(800, 600)
		
		self.add_events(gtk.gdk.POINTER_MOTION_MASK |
						gtk.gdk.BUTTON_PRESS_MASK |
						gtk.gdk.BUTTON_RELEASE_MASK |
						gtk.gdk.ENTER_NOTIFY_MASK |
						gtk.gdk.LEAVE_NOTIFY_MASK |
						gtk.gdk.SCROLL)
		self.connect("button_press_event", self.mouse_button_press)
		self.connect("button_release_event", self.mouse_button_release)
		self.connect("motion-notify-event", self.mouse_move)
		self.connect("scroll-event", self.scroll_event)
		
		self.scene = SceneNode()
		self.reset()
		
		display_mode = (gtk.gdkgl.MODE_RGB | gtk.gdkgl.MODE_DOUBLE | gtk.gdkgl.MODE_DEPTH)
		try:
			glconfig = gtk.gdkgl.Config(mode=display_mode)
		except gtk.gdkgl.NoMatches:
			display_mode &= ~gtk.gdkgl.MODE_SINGLE
			glconfig = gtk.gdkgl.Config(mode = display_mode) 
		self.set_gl_capability(glconfig)
		self.connect("expose_event", self.on_expose_event)
		self.connect("realize", self.on_realize )
		self.connect("configure_event", self.on_configure_event )
	
	def reset(self):
  		self.camera = Camera()
		self.mouse_is_over = False
		self.mouse_position = (0, 0)
		self.mouse_old_position = (0, 0)
		self.rotating = False
		self.panning = False
		self.zooming = False
		self.frame_idx = 0
	
	def mouse_move(self, widget, event):
		if event.type == gtk.gdk.ENTER_NOTIFY:
			self.mouse_is_over = True
		elif event.type == gtk.gdk.LEAVE_NOTIFY:
			self.mouse_is_over = False
		
		rect = self.get_allocation()
		self.mouse_position = (event.x, event.y)
		
		diff = (self.mouse_position[0] - self.mouse_old_position[0],
			   self.mouse_position[1] - self.mouse_old_position[1])
		if self.rotating:
			self.camera.rotate_rel(diff)
			self.draw()
		if self.panning:
			self.camera.translate_rel(diff)
			self.draw()
		if self.zooming:
			self.camera.zoom_rel(diff[1]/100.0 + 1)
			self.draw()
		
		self.mouse_old_position = self.mouse_position
		
	def mouse_button_press(self, widget, event):
		if event.button == 1:
			self.rotating = True
		elif event.button == 2:
			self.panning = True
		elif event.button == 3:
			self.zooming = True
	
	def mouse_button_release(self, widget, event):
		if event.button == 1:
			self.rotating = False
		elif event.button == 2:
			self.panning = False
		elif event.button == 3:
			self.zooming = False
	
	def on_realize(self,*args):
		gldrawable = self.get_gl_drawable()
		glcontext = self.get_gl_context()
		
		if not gldrawable.gl_begin(glcontext):
			return
		gldrawable.gl_end()
		
		glEnable(GL_DEPTH_TEST)
		glDepthFunc(GL_LEQUAL)
		glHint(GL_PERSPECTIVE_CORRECTION_HINT, GL_NICEST)
		
		glLightfv(GL_LIGHT1, GL_AMBIENT, (0.2, 0.2, 0.2, 1.0))
		glLightfv(GL_LIGHT1, GL_DIFFUSE, (0.8, 0.8, 0.8, 1.0))
		glLightfv(GL_LIGHT1, GL_SPECULAR, (0.5, 0.5, 0.5, 1.0))
		glEnable(GL_LIGHT1)
		glEnable(GL_LIGHTING)
		
		self.running = True
		
	def on_configure_event(self, *args):
		# Window resize
		gldrawable = self.get_gl_drawable()
		glcontext = self.get_gl_context()
		gldrawable.gl_begin(glcontext)		
		rect = self.get_allocation()
		glViewport(0, 0, rect.width, rect.height)
		glMatrixMode(GL_PROJECTION)
		glLoadIdentity()
		gluPerspective(70, float(rect.width) / rect.height, 0.001, 10.0)
		glMatrixMode(GL_MODELVIEW)

	def on_expose_event(self, *args):	
		self.draw()
	
	def check_connection(self):
		if self.connection.poll():
			data = self.connection.recv()
			if data[0] == "position":
				print "got data", data
				self.scene[data[1]].position = data[2]
	
	def draw(self):
		self.check_connection()
		
		gldrawable = self.get_gl_drawable()
		glcontext  = self.get_gl_context()
		gldrawable.gl_begin(glcontext)
		glClearColor(0.8,0.8,0.8,0)
		glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
		glLoadIdentity()
		self.camera.set_gl_orientation()
		self.scene.draw()
		
		if self.save_video:
			rect = self.get_allocation()
			width, height = rect.width, rect.height
			pixels = glReadPixelsub(0, 0, width, height, GL_RGB)
			image = Image.fromstring("RGB", (int(width),int(height)), pixels)
			image = image.transpose( Image.FLIP_TOP_BOTTOM)
			image.save("movie/frame%04d.jpg" % self.frame_idx)
		
		gldrawable.swap_buffers()
		gldrawable.gl_end()
		self.frame_idx += 1
		
	def scroll_event(self, widget, event):
		if event.direction == gtk.gdk.SCROLL_UP:
			self.camera.zoom_rel(1.1)
		elif event.direction == gtk.gdk.SCROLL_DOWN:
			self.camera.zoom_rel(0.9)
		self.draw()
	
	def main(self):
		window = gtk.Window()
		window.add(self)
		window.connect("destroy", gtk.main_quit)
		window.show_all()
		gtk.main()

def main():
	window = gtk.Window()
	vw = VisualiseWidget3D()

	window.add(vw)
	window.connect("destroy", gtk.main_quit)
	window.show_all()

	gtk.main()


if __name__ == "__main__":
    main()
